# -*- coding: utf-8 -*-
# @Time    : 2022/2/8 9:43 下午
# @Author  : DIRICHLET
# @Email   : 511827625@qq.com
# @File    : cvae.py
# @Software: PyCharm
# @ Better Late Than Never!
import keras
import pytest
from typing import *

import numpy as np
from tensorflow.keras.utils import to_categorical

def test_1():
    a=0
    assert a==1

def load_data(path)-> Tuple[Tuple[np.ndarray,np.ndarray],Tuple[np.ndarray,np.ndarray]]:

    #统一api
    with np.load(path) as f:
        x_train, y_train = f['x_train'], f['y_train']
        x_test, y_test = f['x_test'], f['y_test']
        return (x_train, y_train), (x_test, y_test)

def dataInit(path):
    (x_train, y_train), (x_test, y_test) = load_data(path=path)
    tfu=lambda d: d.reshape((d.shape[0],-1))/255
    return (tfu(x_train),to_categorical(y_train)),\
           (tfu(x_test),to_categorical(y_test))

def test_cvae_mnist_data():
    import tensorflow as tf  # 导入tensorflow库
    import numpy as np
    #from tensorflow.keras.datasets.mnist import load_data
    #有vpn的时候 可以在线运行代码，获取数据集
    # mnist = tf.keras.datasets.mnist
    # (x_train, y_train), (x_test, y_test) = mnist.load_data()  # 加载数据
    #
    #载入离线数据
    (x_train, y_train), (x_test, y_test) = load_data(path="mnist.npz")  # mnist的本地路径
    print(y_train[0])


def test_cvae_mnist_train():
    from main.algo.ConditionalAutoEncoder import ConditionalAutoEncoder
    import tensorflow as tf

    (x_train, y_train), (x_test, y_test) = dataInit('mnist.npz')
    cvae = ConditionalAutoEncoder(x_train.shape[1], 32, 2, 32, 10)
    cvae.compile(optimizer=tf.keras.optimizers.Adam())
    cvae.fit(x_train, y_train, batch_size=128)
    cvae.save_weights('cvae_model_weight')


def test_cvae_mnist_predict():
    from main.algo.ConditionalAutoEncoder import ConditionalAutoEncoder
    import tensorflow as tf
    import matplotlib.pyplot as plt
    (x_train, y_train), (x_test, y_test) = dataInit('mnist.npz')
    cvae = ConditionalAutoEncoder(x_train.shape[1], 32, 2, 32, 10)
    cvae.load_weights('cave_weight/cvae_model_weight')

    for k in range(10):
        t=np.array([[0 for kk in range(10)]])
        t[0][k]=1
        p=cvae(x_train[0:1],t,'predict')
        print(t)
        plt.figure(figsize=(1.0 * 1, 1 * 2))
        for i in range(1):
            # 原始图像
            plt.subplot(2, 1, i + 1)
            plt.imshow(np.reshape(x_train[0]*255, (28, 28)), cmap='gray')
            plt.axis('off')

            # 根据标签成成模拟数据
            plt.subplot(2, 1, i + 1 + 1)
            plt.imshow(np.reshape(p*255, (28, 28)), cmap='gray')
            plt.axis('off')
        plt.show()





